import cv2
import numpy as np

cap = cv2.VideoCapture("../video/vtest.avi")
ret, frame1 = cap.read()
prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
hsv = np.zeros_like(frame1)
hsv[..., 1] = 255   #...表示最后一维前面的所有维度，最为一维为1，就表示将饱和度设为255
while (1):
    ret, frame2 = cap.read()
    next = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
    #利用前一帧和当前帧进行光流计算，计算得到的光流也是和图像一样的尺寸，元素为对应像素的光流方向矢量
    flow = cv2.calcOpticalFlowFarneback(prvs, next, None, 0.5, 3, 15, 3, 5, 1.2, 0)
    #矢量转换到极坐标，得到幅值和角度
    mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
    hsv[..., 0] = ang * 180 / np.pi / 2     #色调为对应像素光流方向的角度到0到180之间的映射
    hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX) #明度设为对应像素光流方向的幅度到0到255之间的映射
    bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)  #再转换到BGR颜色空间
    cv2.imshow('frame2', bgr)
    cv2.imshow('origin',frame2)
    k = cv2.waitKey(30) & 0xff
    if k == 27: #esc退出
        break
    elif k == ord('s'):
        cv2.imwrite('opticalfb.png', frame2)
        cv2.imwrite('opticalhsv.png', bgr)
    prvs = next
cap.release()
cv2.destroyAllWindows()